#include <cuda.h>
#include <stdio.h>



#define OLAMDA (0)


extern "C"  void miniTV_gpu(float *x_recon,float *x0, int width, int height, float lamda,int iterNum);
cudaArray* initCudaArrayTexture4Grad(texture<float, 3,cudaReadModeElementType> &texArray,int width,int height, int length, float *HostData);
__global__ void kernelGradDesSlice(float * Dx,int size,float mu,int indexZ);



  //texture<float, 3,cudaReadModeElementType> texPriVolume;
       texture<float, 3,cudaReadModeElementType> texVolumeART;
       texture<float, 3,cudaReadModeElementType> texVolumeGrad;
	   
	   
	   

void checkCUDAError(const char *msg)
{

    cudaError_t err = cudaGetLastError();
    if( cudaSuccess != err) 
    {
        fprintf(stderr, "Cuda error: %s: %s.\n", msg, cudaGetErrorString( err) );
        exit(EXIT_FAILURE);
    }                         
}

void UpdateVolume(float *tempDx, float *x_recon, int size, float t)
{

  for(int index =0;index<size;index++)
  {
    x_recon[index] = t*tempDx[index] + x_recon[index]; 
  }
}


__global__ void kernalGetObjFun(float *tmpResutls,float *Dx,int size,float t,int indexZ,float mu)
{


	
	  int z= (indexZ);
	int x = (threadIdx.x) + (blockIdx.x)* 16;
	int y = (threadIdx.y) + (blockIdx.y)* 16; 

    float sumTV = 0;


    float GH = 0;
    float GV = 0;
	//float GHDIF = 0;
	//float GVDIF = 0;

    if(x> (size -1) || y> (size -1)) return;
	if( ( y<(size -1) ) && ( x<(size -1) ) )
	{
		//|x|TV
		
		GH = - tex3D(texVolumeGrad,x,y,z) + tex3D(texVolumeGrad,x  ,y+1,z);
		GV = - tex3D(texVolumeGrad,x,y,z) + tex3D(texVolumeGrad,x+1,y  ,z);
		
		if(t!=0)
		{
            GH += t * (-Dx[x + y * size ] + Dx[ x    + (y +1 )* size ]);   
            GV += t * (-Dx[x + y * size ] + Dx[(x+1) +  y     * size ]);
		}

		sumTV += sqrtf(GH*GH + GV*GV);
	
		
	}

	tmpResutls[x + y*size] = sumTV * mu ;
     
}

float ObjFun(float *Dx,int size,float t,int indexZ,float mu,float *volumeART,float *currentVolume,float * hostDx)
{
     float *DtmpResults = NULL;
	 int numVoxels =  size * size;
	 cudaMalloc((void **) &(DtmpResults), numVoxels*sizeof(float)); checkCUDAError("malloc objFun tmpResults");
     
	 float finalResults =0;
	 

	dim3 grid_dbp(size/16, size/16);
	dim3 testgrid(16, 16);
	   
	 //kernalGetObjFun<<<size,size>>>(DtmpResults,Dx,size,t,indexZ,mu);
	 kernalGetObjFun<<<grid_dbp,testgrid>>>(DtmpResults,Dx,size,t,indexZ,mu);

	float *HtmpResults = (float *)malloc(sizeof(float)*size*size);
	cudaMemcpy(HtmpResults,DtmpResults,sizeof(float)*size*size,cudaMemcpyDeviceToHost);

	 for(int i =0;i<size*size;i++)
	 {
		 finalResults += HtmpResults[i];
		 finalResults += ( volumeART[i+ indexZ * numVoxels] - currentVolume[i+ indexZ * numVoxels] - hostDx[i] )
                  		 * ( volumeART[i+ indexZ * numVoxels] - currentVolume[i+ indexZ * numVoxels] - hostDx[i] );
	 }


	 free(HtmpResults);
	 cudaFree(DtmpResults);

	  return finalResults;
}

float norm(float *Dx,int size)
{
   float *hostDx = (float *)malloc(size*sizeof(float));
   cudaMemcpy(hostDx,Dx,sizeof(float)*size,cudaMemcpyDeviceToHost);
   
   float normValue = 0.0f;
   for(int i=0;i<size;i++)
   {
     normValue = normValue + hostDx[i];
   }
   
   free(hostDx);
   
   return normValue;
}

//
float searchLine(float *Dx,int size,int indexZ,float mu,float *volumeART,float *currentVolume, float * hostDx)
{
    float alpha = 0.01f;  
	float beta = 0.5f;  
	int backiter = 0;
	float t = 1.0f ; 
    int maxIter = 30;	


	float f0 = ObjFun(Dx,size,0.0f,indexZ,mu,volumeART,currentVolume,hostDx); //f0
	float dtg = norm(Dx,size*size);
	float fl = 0.0f;
	float fr = 0.0f;
	while(backiter<maxIter)
	{
        backiter++;
		fl = ObjFun(Dx,size,t,indexZ,mu,volumeART,currentVolume,hostDx);
        fr = f0 + alpha*t*dtg;
		if(fr>=fl) break;
        
		//printf("search line iter ====== %d \n",backiter);

		t = beta*t;


	}
    //if(backiter == maxIter) t = 0.0f;
    return t;

}

extern "C"  void miniTV_gpu(float *x_recon,float *x0, int width, int height, float lamda,int iterNum)
{
       int numVoxels = width * height;
	   int length =1;
	   
	   memcpy((void *)x_recon,(void*) x0,sizeof(float)*width*height);
	   
     
	   
	
	   
	
	   
	   cudaArray *cudaVolumeART = initCudaArrayTexture4Grad(texVolumeART,width,height,length,x0);
       cudaArray *cudaVolumeGrad = initCudaArrayTexture4Grad(texVolumeGrad,width,height,length,x0);
	   
	  
	    


	   
	   //dx
       float *Dx;
	  cudaMalloc((void **) &(Dx), numVoxels*sizeof(float)); checkCUDAError("malloc gradient");
	   cudaMemset(Dx, 0, numVoxels*sizeof(float));  checkCUDAError("cudaMemset ");
	     
	    //
	   dim3 grid_dbp(width/16, height/16);
	   dim3 testgrid(16, 16);
	   
	    float t = 1.0f;

	

	 
	  for(int tvIter =0;tvIter<iterNum;tvIter++)
	  {


		  //get dx
		  t = 1.0f;
		   cudaMemset(Dx, 0, width*height*sizeof(float));  checkCUDAError("cudaMemset1 ");
		   kernelGradDesSlice<<<grid_dbp,testgrid>>>(Dx,width,lamda,1);checkCUDAError("kernelGradDesSlice");
		
			  
            
		   float *tempDx = (float *) malloc(sizeof(float) * width * height);
		      cudaMemcpy(tempDx,Dx,width*height*sizeof(float), cudaMemcpyDeviceToHost); 
			  checkCUDAError("cudaMemcpy  tempDx");
			  
			  t  = searchLine(Dx ,width,0,lamda,x0,x_recon,tempDx);
			
			
         
			  
			   //update volume
              UpdateVolume(tempDx, x_recon, width*height, t);
			   
			   free(tempDx);
			   
			   //update in texture
		  cudaMemcpy3DParms copyParams = {0};
	      cudaExtent volumeSize = make_cudaExtent(width, height, 1);
		  copyParams.srcPtr = make_cudaPitchedPtr((void*)(x_recon), width*sizeof(float),width,width);
          
		  copyParams.dstArray = cudaVolumeGrad;
	      copyParams.extent = volumeSize;
	      copyParams.kind = cudaMemcpyHostToDevice;

		  cudaMemcpy3D(&copyParams);
	      checkCUDAError("update volumeGrad in texture");
			 
          
			  
	  }//end for TV iter



	   
	   
      //free(VolumeART);
      
	  cudaUnbindTexture(texVolumeART);
	  cudaUnbindTexture(texVolumeGrad);


	  cudaFreeArray(cudaVolumeART);
	  cudaFreeArray(cudaVolumeGrad);
	  
	  cudaFree(Dx);
	   

  
}

__global__ void kernelGradDesSlice(float * Dx,int size,float mu,int indexZ)
{
  
    int z= (indexZ);
	int x = (threadIdx.x) + (blockIdx.x)* 16;
	int y = (threadIdx.y) + (blockIdx.y)* 16; 

	Dx[x+y*size] = 0.0f;


  
	if(x> (size -1) || y> (size -1)) return;

	
	float sum=0;
	
   //grad volume_TV
	float Dh = tex3D(texVolumeGrad,x+1,y,z) - tex3D(texVolumeGrad,x,y,z);
	float Dv = tex3D(texVolumeGrad,x,y+1,z) - tex3D(texVolumeGrad,x,y,z);
	

 

    float Dyh = 0.0f;
	float Dyv = 0.0f;

	float Dxh = 0.0f;
	float Dxv = 0.0f;

	if(y>0&&x>0)
	{

	     Dyh = tex3D(texVolumeGrad,x+1,y-1,z) - tex3D(texVolumeGrad,x,y-1,z);
	     Dyv = tex3D(texVolumeGrad,x,y,z) - tex3D(texVolumeGrad,x,y-1,z);

	     Dxh = tex3D(texVolumeGrad,x,y,z) - tex3D(texVolumeGrad,x-1,y,z);
	     Dxv = tex3D(texVolumeGrad,x-1,y+1,z) - tex3D(texVolumeGrad,x-1,y,z);

	}
	

          
	  float smallNum = 0.00000000000000001f;
	  float s = sqrt(Dh*Dh + Dv*Dv);
	  if(s<smallNum) s = smallNum;
	  sum += (1 - OLAMDA)*(Dh + Dv)/s;
	  
	  

	   //
	  s = sqrt(Dyh*Dyh + Dyv*Dyv);
	  if(s<smallNum) s = smallNum;
	  sum += (Dyv)*(OLAMDA -1.0f )/s;

	  //
	  s = sqrt(Dxh*Dxh + Dxv*Dxv);
	  if(s<smallNum) s = smallNum;
	  sum += (Dxh)*(OLAMDA -1.0f )/s;
      
	  Dx[ x + y*size] += mu * sum;
	  
	  //texVolumeART
	  Dx[ x + y*size ] += tex3D(texVolumeART,x,y,z) - tex3D(texVolumeGrad,x,y,z); 
	  

   
}


 cudaArray* initCudaArrayTexture4Grad(texture<float, 3,cudaReadModeElementType> &texArray,int width,int height, int length, float *HostData)
{
       
       cudaArray *cudaV;
	   cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<float>();  
       cudaExtent extent;     
       extent.width= width;     
       extent.height=height;      
       extent.depth=length;   
       
	   //alloc cuda array
       cudaMalloc3DArray(&cudaV,&channelDesc,extent);
	   checkCUDAError("cudaMalloc3DArray cudaPriVolume");

	   //bind to texture and init
	   texArray.addressMode[0] = cudaAddressModeClamp;
	   texArray.addressMode[1] = cudaAddressModeClamp;
	   texArray.addressMode[2] = cudaAddressModeClamp;
	   texArray.filterMode = cudaFilterModePoint;
	   texArray.normalized = 0; 
       cudaBindTextureToArray(texArray, cudaV,channelDesc); 
	   checkCUDAError("cudaBindTextureToArray texVolume");
	   
	   cudaMemcpy3DParms copyParams = {0};
	   copyParams.srcPtr = make_cudaPitchedPtr((void*)(HostData), 
						width*sizeof(float), 
						height, 
						width);
	   copyParams.dstArray = cudaV;
	   copyParams.extent = extent;
	   copyParams.kind = cudaMemcpyHostToDevice;

	   cudaMemcpy3D(&copyParams);
	   checkCUDAError("cudaMemcpy3D cudaDynaIterCopyVolumeHostToDeviceArray");

	   return cudaV;

}